#include <stdio.h>
#include <pthread.h>
#include <time.h>
#include "../../src/queue/lockfree_queue.h"

#define NUM_PRODUCERS 4
#define NUM_CONSUMERS 4
#define ITEMS_PER_PRODUCER 100000

queue_t queue;
atomic_int produced = 0;
atomic_int consumed = 0;

void* producer(void* arg) {
    for (int i = 0; i < ITEMS_PER_PRODUCER; i++) {
        int* value = malloc(sizeof(int));
        *value = i;
        while (!queue_enqueue(&queue, value)) {
            // 如果入队失败，重试
            sched_yield();
        }
        atomic_fetch_add(&produced, 1, memory_order_relaxed);
    }
    return NULL;
}

void* consumer(void* arg) {
    while (1) {
        void* data = queue_dequeue(&queue);
        if (data) {
            free(data);
            atomic_fetch_add(&consumed, 1, memory_order_relaxed);
        } else if (atomic_load(&produced, memory_order_relaxed) == 
                  NUM_PRODUCERS * ITEMS_PER_PRODUCER &&
                  atomic_load(&consumed, memory_order_relaxed) == 
                  atomic_load(&produced, memory_order_relaxed)) {
            break;
        } else {
            sched_yield();
        }
    }
    return NULL;
}

int main() {
    pthread_t producers[NUM_PRODUCERS];
    pthread_t consumers[NUM_CONSUMERS];
    struct timespec start, end;
    
    queue_init(&queue);
    clock_gettime(CLOCK_MONOTONIC, &start);
    
    // 创建生产者线程
    for (int i = 0; i < NUM_PRODUCERS; i++) {
        pthread_create(&producers[i], NULL, producer, NULL);
    }
    
    // 创建消费者线程
    for (int i = 0; i < NUM_CONSUMERS; i++) {
        pthread_create(&consumers[i], NULL, consumer, NULL);
    }
    
    // 等待所有线程完成
    for (int i = 0; i < NUM_PRODUCERS; i++) {
        pthread_join(producers[i], NULL);
    }
    for (int i = 0; i < NUM_CONSUMERS; i++) {
        pthread_join(consumers[i], NULL);
    }
    
    clock_gettime(CLOCK_MONOTONIC, &end);
    
    double time_spent = (end.tv_sec - start.tv_sec) +
                       (end.tv_nsec - start.tv_nsec) / 1e9;
    
    printf("Time spent: %.2f seconds\n", time_spent);
    printf("Operations per second: %.2f\n", 
           (NUM_PRODUCERS * ITEMS_PER_PRODUCER) / time_spent);
    printf("Total items produced: %d\n", 
           atomic_load(&produced, memory_order_relaxed));
    printf("Total items consumed: %d\n", 
           atomic_load(&consumed, memory_order_relaxed));
    
    return 0;
} 